
import cv2
import numpy as np
import matplotlib.pyplot as plt

# Get occurrence of each gray 
def GetHistogram(image, width, height, depth):
    '''Get Histogram statistics information
    '''
    grayLevel = 2**depth
    a = [0]*(grayLevel)
    for i in range(height):
        for j in range(width):
            iGray = int(image[i,j])
            a[iGray] = a[iGray] + 1
    return a


# Calculator new intensity value
def EqualizeHistorgram(originalHist, width, height, depth):
    '''Computer equalization equation. Return a gray level mapping table.
    '''
    grayLevel = 2**depth
    histMappingTable = [0]*(grayLevel)
    MN = width*height
    tmpSum = 0
    for i in range(grayLevel):
        tmpSum += originalHist[i]
        histMappingTable[i] = round((grayLevel-1)*tmpSum*1.0/MN)
        # if i  <= 10:
        #     print i
        #     print originalHist[i]
        #     print histMappingTable[i] 
    return histMappingTable


# Apply equalized histogram to original image
def ApplyEqualization(targetImage, originalImage, width, height, histMappingTable):
    for i in range(height):
        for j in range(width):
            iGray = int(originalImage[i,j])
            #if (i <= 100) and (j <= 100):
            #    print histMappingTable[iGray]
            targetImage[i,j] = histMappingTable[iGray]
    return targetImage



def EqualizeImage(rawImg, displayHistogram = True):
    '''Compute the histogram of input image.
        rawImg is the raw input image. e.g. cv2.imread('image path', cv2.CV_LOAD_IMAGE_GRAYSCALE)
        displayHistogram control if display histogram plot
    '''
    height, width = rawImg.shape
    depth = rawImg.dtype.itemsize * 8 # get depth level, 8 bit or 16 bit
    # Get occurrence count of each gray level
    rawHist = GetHistogram(rawImg, width, height, depth)
    # Apply equalization equation to raw histogram
    histMapping = EqualizeHistorgram(rawHist, width, height, depth)
    # init result image
    resultImage = rawImg.copy() 
    # Transform intensity value of original image
    resultImage = ApplyEqualization(resultImage, rawImg, width, height, histMapping)

    if displayHistogram:
        plt.figure()
        plt.plot(rawHist)

        plt.figure()
        plt.plot(histMapping)

        plt.figure()
        resultHist = GetHistogram(resultImage, width, height, depth)
        plt.plot(resultHist)

    return resultImage



if __name__ == '__main__':
    rawFig1 = cv2.imread('./img/Fig1.jpg', cv2.CV_LOAD_IMAGE_GRAYSCALE)
    cv2.namedWindow('Original Fig1', cv2.CV_WINDOW_AUTOSIZE)
    cv2.imshow( "Original Fig1", rawFig1)

    ######
    resultImage = EqualizeImage(rawFig1)
    ######

    cv2.namedWindow('Equalized Fig1', cv2.CV_WINDOW_AUTOSIZE)
    cv2.imshow( "Equalized Fig1", resultImage)

    #equ = cv2.equalizeHist(rawFig1)  
    #cv2.imshow('equ',equ)  

    rawFig2 = cv2.imread('./img/Fig2.jpg', cv2.CV_LOAD_IMAGE_GRAYSCALE)
    cv2.namedWindow('Original Fig2', cv2.CV_WINDOW_AUTOSIZE)
    cv2.imshow( "Original Fig2", rawFig2)

    ######
    resultImage2 = EqualizeImage(rawFig2)
    ######

    cv2.namedWindow('Equalized Fig2', cv2.CV_WINDOW_AUTOSIZE)
    cv2.imshow( "Equalized Fig2", resultImage2)

    plt.show()
    cv2.waitKey(0)
    cv2.destroyAllWindows()